/*
 * simple_gm.cc
 *
 *  Created on: Sep 9, 2011
 *      Author: tqlong
 */

#include "simple_gm.h"
#include <algorithm>
#include <iostream>
#include <fstream>
#include <boost/foreach.hpp>

Edge::Edge(int i, int j, bool unique)
{
  if (!unique)
  {
    i_ = std::min(i, j);
    j_ = std::max(i, j);
  }
  else
  {
    i_ = i;
    j_ = j;
  }
}

bool Edge::EdgeCompare::operator()(const Edge& lhs, const Edge& rhs) const
{
  return (lhs.i_ < rhs.i_) || (lhs.i_ == rhs.i_ && lhs.j_ < rhs.j_);
}

void SimpleGM::load(const char* fileName)
{
  std::ifstream file(fileName);
  if (!file.good())
  {
    std::cerr << "Error reading file: " << fileName << std::endl;
    exit(1);
  }

  // get the number of nodes
  int n;
  file >> n;
  //std::cout << "Load GM: n=" << n << std::endl;
  // get node potentials
  node_potentials_.clear();
  for (int i = 0; i < n; i++)
  {
    int card;
    file >> card;
    arma::vec v(card);
    for (int j = 0; j < card; j++)
    {
      file >> v[j];
    }

    int idx = addNode(card);
    addPotential(idx, v);
    //std::cout << "Load GM: card(" << idx << ")=" << cardinality_[idx] << std::endl;
    //std::cout << "Load GM: pot(" << idx << ")=" << std::endl << node_potentials_[idx];
  }

  int ne;
  file >> ne;
  //std::cout << "Load GM: ne=" << ne << std::endl;
  edge_potentials_.clear();
  for (int k = 0; k < ne; k++)
  {
    int i, j;
    file >> i >> j;
    arma::mat table(cardinality_[i], cardinality_[j]);
    for (int vj = 0; vj < cardinality_[j]; vj++)
      for (int vi = 0; vi < cardinality_[i]; vi++)
      {
        double v;
        file >> v;
        table.at(vi, vj) = v;
      }

    addEdge(i, j);
    addPotential(i, j, table);
    //std::cout << "Load GM: pot(" << e.i_ << "," << e.j_ << ")="  << std::endl << edge_potentials_[e];
  }

  file.close();
}

int SimpleGM::addNode(int cardinality)
{
  int i = cardinality_.size();
  cardinality_[i] = cardinality;
  return i;
}

void SimpleGM::addPotential(int i, const arma::vec& v)
{
  if (node_potentials_.find(i) == node_potentials_.end())  // no potential registered
    node_potentials_[i] = v;
  else // there is a registered potential for node i
    node_potentials_[i] += v;
}

Edge SimpleGM::addEdge(int i, int j)
{
  Edge e(i, j);
  adj_edges_[i].push_back(e);
  adj_edges_[j].push_back(e);
  return e;
}

void SimpleGM::addPotential(int i, int j, const arma::mat& v)
{
  Edge e(i, j);
  arma::mat tmp = v;
  if (i >= j) tmp = arma::trans(v);
  if (edge_potentials_.find(e) == edge_potentials_.end())
    edge_potentials_[e] = v;
  else
    edge_potentials_[e] += v;
}

double SimpleGM::pseudo_probability(const imap_type& a, bool isLog)
{
  if (!isCompleteAssignment(a))
  {
    std::cerr << "the assignment is not compelete: " << __FILE__ << ":" << __LINE__ << std::endl;
    return 0;
  }

  double s = 0;

  typedef const std::pair<Node, arma::vec>& pair_node_t;
  BOOST_FOREACH(pair_node_t p, node_potentials_)
  {
    const Node& i = p.first;
    const arma::vec& v = p.second;
    int vi = a.find(i)->second;
    double pot = v.at(vi);
    s += pot;
  }

  typedef const std::pair<Edge, arma::mat>& pair_edge_t;
  BOOST_FOREACH(pair_edge_t p, edge_potentials_)
  {
    const Edge& e = p.first;
    const arma::mat& v = p.second;
    int vi = a.find(e.i_)->second;
    int vj = a.find(e.j_)->second;
    double pot = v.at(vi, vj);
    s += pot;
  }

  return isLog ? s : exp(s);
}

double SimpleGM::partition_function(bool isLog)
{
  std::vector<int> vi(n(), 0);
  assignment_type a;

  double s = log(0);
  while (1)
  {
    for (int i = 0; i < n(); i++) a[i] = vi[i];
    double pHat = pseudo_probability(a); // log pseudo probability
    s = logadd(s, pHat);

    // information
    //std::cout << "log_pHat(";
    //for (int i = 0; i < n; i++) std::cout << vi[i] << ",";
    //std::cout << ") = " << pHat << " " << exp(pHat) << std::endl;

    // generate next assignment
    if (!next_assignment(vi)) break;
  }

  return isLog ? s : exp(s);
}

void SimpleGM::naive_marginals()
{
  init_marginals();
  // Naive computation
  std::vector<int> vi(n(), 0);
  assignment_type a;

  typedef std::pair<const Node, arma::vec>& pair_node_t;
  typedef std::pair<const Edge, arma::mat>& pair_edge_t;

  double s = log(0);
  while (1)
  {
    for (int i = 0; i < n(); i++) a[i] = vi[i];
    double pHat = pseudo_probability(a); // log pseudo probability
    s = logadd(s, pHat);

    // node accumulation
    BOOST_FOREACH(pair_node_t p, node_marginals_)
    {
      Node i = p.first;
      arma::vec& v = p.second;
      double tmp = logadd(v.at(vi[i]), pHat);
      v.at(vi[i]) = tmp;
    }

    // edge accumulation
    BOOST_FOREACH(pair_edge_t p, edge_marginals_)
    {
      Node i = p.first.i_, j = p.first.j_;
      arma::mat& v = p.second;
      double tmp = logadd(v.at(vi[i], vi[j]), pHat);
      v.at(vi[i], vi[j]) = tmp;
    }

    // generate next assignment
    if (!next_assignment(vi)) break;
  }
}

SimpleGM::assignment_type SimpleGM::naive_map()
{
  // Naive computation
  std::vector<int> vi(n(), 0);
  assignment_type a;

  typedef std::pair<const Node, arma::vec>& pair_node_t;
  typedef std::pair<const Edge, arma::mat>& pair_edge_t;

  double max_p = log(0);
  assignment_type max_a;

  while (1)
  {
    for (int i = 0; i < n(); i++) a[i] = vi[i];
    double pHat = pseudo_probability(a); // log pseudo probability
    if (pHat > max_p)
    {
      max_p = pHat;
      max_a = a;
    }

    // information
    //std::cout << "log_pHat(";
    //for (int i = 0; i < n; i++) std::cout << vi[i] << ",";
    //std::cout << ") = " << pHat << " " << exp(pHat) << std::endl;

    // generate next assignment
    if (!next_assignment(vi)) break;
  }
  return max_a;
}

// Reduce a graphical model based on evidence
SimpleGM SimpleGM::reduce(const assignment_type& a)
{
  imap_type old2new, new2old;
  SimpleGM reduced;

  // process non-evidence nodes
  typedef const std::pair<const Node, arma::vec>& pair_node_t;
  BOOST_FOREACH(pair_node_t p, node_potentials_)
  {
    int oldNode = p.first;
    if (a.find(oldNode) != a.end()) continue;
    int newNode = reduced.addNode(cardinality_[oldNode]);
    reduced.addPotential(newNode, p.second);
    old2new[oldNode] = newNode;
    new2old[newNode] = oldNode;
  }

  // process edge_potentials
  typedef const std::pair<const Edge, arma::mat>& pair_edge_t;
  BOOST_FOREACH(pair_edge_t p, edge_potentials_)
  {
    int oldi = p.first.i_, oldj = p.first.j_;
    const arma::mat& v = p.second;
    imap_type::const_iterator it = a.find(oldi);
    if (it != a.end())
    {
      int vali = it->second;
      if (a.find(oldj) == a.end())  // oldi is known, oldj is not known
      {
        int newj = old2new[oldj];
        arma::vec newv = arma::trans(v.row(vali));
        reduced.addPotential(newj, newv);
      }
      // else {}                    // both are known so we should not do anything
    }
    else
    {
      imap_type::const_iterator it = a.find(oldj);
      if (it == a.end())  // oldi and oldj are not known
      {
        int newi = old2new[oldi], newj = old2new[oldj];
        reduced.addEdge(newi, newj);
        reduced.addPotential(newi, newj, v);
      }
      else                          // old i is not known, oldj is known
      {
        int valj = it->second;
        int newi = old2new[oldi];
        arma::vec newv = v.col(valj);
        reduced.addPotential(newi, newv);
      }
    }
  }
  reduced.old2new_ = old2new;
  reduced.new2old_ = new2old;
  return reduced;
}

// Information printing
void SimpleGM::print_params() const
{
  typedef const std::pair<const Node, arma::vec>& pair_node_t;
  typedef const std::pair<const Edge, arma::mat>& pair_edge_t;
  // node
  BOOST_FOREACH(pair_node_t p, node_potentials_)
  {
    Node i = p.first;
    const arma::vec& v = p.second;
    std::cout << "cardinality(" << i << ")=" << cardinality_.find(i)->second << std::endl;
    std::cout << "log_potential(" << i << ")=" << std::endl << v;
  }

  // edge
  BOOST_FOREACH(pair_edge_t p, edge_potentials_)
  {
    Node i = p.first.i_, j = p.first.j_;
    const arma::mat& v = p.second;
    std::cout << "log_potential(" << i << "," << j << ")=" << std::endl << v;
  }
}

void SimpleGM::print_marginals() const
{
  typedef const std::pair<const Node, arma::vec>& pair_node_t;
  typedef const std::pair<const Edge, arma::mat>& pair_edge_t;
  // node marginals
  BOOST_FOREACH(pair_node_t p, node_marginals_)
  {
    Node i = p.first;
    const arma::vec& v = p.second;
    std::cout << "marginal(" << i << ")=" << logsum(v) << std::endl << v;
  }

  // edge marginals
  BOOST_FOREACH(pair_edge_t p, edge_marginals_)
  {
    Node i = p.first.i_, j = p.first.j_;
    const arma::mat& v = p.second;
    std::cout << "marginal(" << i << "," << j << ")=" << logsum(v) << std::endl << v;
  }
}

/** Protected members */
bool SimpleGM::isCompleteAssignment(const imap_type& a) const
{
  int n = cardinality_.size();
  for (int i = 0; i < n; i++)
  {
    imap_type::const_iterator it = a.find(i);
    if (it == a.end())
      return false; // value not available
    if (it->second >= cardinality_.find(i)->second)
      return false; // value out of bound
  }
  return true;
}

double SimpleGM::nodePotential(int i, int value) const
{
  double pot = node_potentials_.find(i)->second.at(value);
  return pot;
}

void SimpleGM::init_marginals()
{
  // Init marginals to zeros
  node_marginals_ = node_potentials_;
  edge_marginals_ = edge_potentials_;
  typedef std::pair<const Node, arma::vec>& pair_node_t;
  BOOST_FOREACH(pair_node_t p, node_marginals_)
  {
    arma::vec& v = p.second;
    v.fill(log(0));
    //std::cout << "init marginal(" << p.first << ")=" << std::endl << node_marginals_[p.first];
  }

  typedef std::pair<const Edge, arma::mat>& pair_edge_t;
  BOOST_FOREACH(pair_edge_t p, edge_marginals_)
  {
    arma::mat& v = p.second;
    v.fill(log(0));
    //std::cout << "init marginal(" << p.first.i_ << "," << p.first.j_ << ")=" << std::endl << edge_marginals_[p.first];
  }
}

bool SimpleGM::next_assignment(std::vector<int>& vi)
{
  int n = vi.size();
  int i = 0;
  while (i < n && vi[i] + 1 == cardinality_[i]) i++;
  if (i == n) return false; // no more assignment
  vi[i--]++;
  while (i >= 0) vi[i--] = 0;
  return true;
}

double SimpleGM::edge_potentials(int i, int j, int vi, int vj) const
{
  Edge e(i, j);
  const arma::mat& v = edge_potentials_.find(e)->second;
  return (i < j) ? v.at(vi, vj) : v.at(vj, vi);
}

/** Log space computation */
double logadd(double a, double b)
{
  if (a == -std::numeric_limits<double>::infinity())
    return b;
  else if (b == -std::numeric_limits<double>::infinity())
    return a;
  else if (a < b)
    return b+log1p(exp(a-b));
  else // a >= b
    return a+log1p(exp(b-a));
}

double logsum(const arma::vec& v)
{
  double s = log(0);
  for (unsigned int i = 0; i < v.n_elem; i++)
    s = logadd(s, v.at(i));
  return s;
}

double logsum(const arma::mat& v)
{
  double s = log(0);
  for (unsigned int i = 0; i < v.n_rows; i++)
    for (unsigned int j = 0; j < v.n_cols; j++)
      s = logadd(s, v.at(i,j));
  return s;
}
